import sys
import time
import serial
import string
import binascii
import pyqtgraph as pg
import array
import threading
import numpy as np
from queue import Queue
from scipy import signal

## 此程序文件主要用于脑电特定信号采集 ##

i = 0
j=0
bb,aa = signal.butter(6, [14 / 256 * 2, 80 / 256 * 2], btype='bandpass')


#脑电波串口读取
def Serial():
    global i 
    while(True):
        if mSerial.inWaiting() > 0:
            datat0=mSerial.readline()
            datat=mSerial.readline()
            data=datat.decode().split(",")
            #Meditation  = int(data[0])+int(data[1])
            while(data[0]=='\n'):
                datat0=mSerial.readline()
                datat=mSerial.readline()
                data=datat.decode().split(",")
            if i < historyLength:
                ab_data[i]=int(data[0])+int(data[1])
                i=i+1
            else:
                ab_data[:-1] = ab_data[1:]
                ab_data[i-1]=int(data[0])+int(data[1])
        
#脑电波实时绘图
def plotData():
    global i,j
    curve1.setData(ab_data,pen = 'b')
    
    if i>1000:
        gamma_data=signal.filtfilt(bb,aa,ab_data)
        gamma_data = np.apply_along_axis(lambda x: np.abs(signal.hilbert(x)), 0,gamma_data)
        curve2.setData(gamma_data,pen = 'm')

        gamma_average = np.mean(gamma_data[i-1000:i])
        if (j>200 and gamma_data[i-96]>gamma_average*2 and gamma_data[i-96]>20):
            print('皱眉')
            temp_data=gamma_data[i-129:i-1]
            curve3.setData(temp_data,pen = 'r')            
            j=0

            print('Array:\n', temp_data)
            file = open("temp_data.txt", "a")
            #content = str(temp_data)
            content = ' '.join(str(elem) for elem in temp_data)
            content=content+'\n'
            file.write(content)
            file.close()

        j=j+1



if __name__ == "__main__":
    app = pg.mkQApp()
    win=pg.GraphicsLayoutWidget(show=True)
    win.setWindowTitle(u"脑电波实时检测")
    win.resize(1000,700) 
    EEG1_data = array.array('i')
    historyLength = 3000    #脑电波数据长度
 
    ab_data = np.zeros(historyLength).__array__('d')
    abc_data = np.zeros(historyLength).__array__('d')
    gamma_data = np.zeros(historyLength).__array__('d')
    temp_data = np.zeros(historyLength).__array__('d')
    

    EEG1 = win.addPlot(left = 'y',buttom = 'x',title = "total")
    EEG1.setRange(xRange = [0,historyLength],yRange=[0,2500],padding = 0)
    curve1 = EEG1.plot()
    curve1.setData(abc_data)
    win.nextRow()

    EEG2 = win.addPlot(left = 'y',buttom = 'x',title = "gamma")
    EEG2.setRange(xRange = [0,historyLength],yRange=[0,300],padding = 0)
    curve2 = EEG2.plot()
    curve2.setData(gamma_data)
    win.nextRow()

    EEG3 = win.addPlot(left = 'y',buttom = 'x',title = "temp")
    EEG3.setRange(xRange = [0,historyLength],yRange=[0,300],padding = 0)
    curve3 = EEG3.plot()
    curve3.setData(temp_data)

    portx = 'COM4'  #这里的端口根据脑环接入电脑的实际虚拟端口进行更正
    bps = 115200
    mSerial = serial.Serial(portx, int(bps))    #打开串口
    if(mSerial.isOpen()):
        print('serial is OK')
    else:
        print("serial is ERROR")
        mSerial.close()  # 关闭端口
    th1 = threading.Thread(target=Serial)
    th1.start()
    timer = pg.QtCore.QTimer()
    timer.timeout.connect(plotData)
    timer.start(4)
    sys.exit(app.exec())
